package com.at.bigdata.spark.core.acc

import org.apache.spark.util.{AccumulatorV2, LongAccumulator}
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * 依赖关系
 *
 * @author cdhuangchao3
 * @date 2023/3/6 8:40 PM
 */
object Spark04_Acc_WordCount {

  def main(args: Array[String]): Unit = {
    val sparConf = new SparkConf().setMaster("local").setAppName("WordCount")
    val sc = new SparkContext(sparConf)

    val rdd = sc.makeRDD(List("hello", "spark", "hello"))

    // 累加器：WordCount
    // 创建累加器对象
    val wcAcc = new MyAccumulator()
    sc.register(wcAcc, "wordCountAcc")
    rdd.foreach(
      word => {
        // 数据的累加
        wcAcc.add(word)
      }
    )
    println("sum=" + wcAcc.value)

    sc.stop()
  }

  /**
   * 自定义累加器
   * 1.继承AccumulatorV2，定义泛型
   *  IN：累加器输入的数据类型
   *  OUT：累加器返回的数据类型
   * 2.重写方法
   */
  class MyAccumulator extends AccumulatorV2[String, mutable.Map[String, Long]] {
    private val _count = mutable.Map[String, Long]()

    // 判断是否初始状态
    override def isZero: Boolean = {
      _count.isEmpty
    }

    override def copy(): AccumulatorV2[String, mutable.Map[String, Long]] = {
      new MyAccumulator
    }

    override def reset(): Unit = {
      _count.clear()
    }

    /**
     * 获取累加器需要计算的值
     * @param v
     */
    override def add(v: String): Unit = {
      _count.update(v, _count.getOrElse(v, 0L) + 1L)
    }

    /**
     * Driver合并多个累加器
     * @param other
     */
    override def merge(other: AccumulatorV2[String, mutable.Map[String, Long]]): Unit = {
      val map1 = this._count
      val map2 = other.value
      map2.foreach{
        case (word, count) => {
          val newCount = map1.getOrElse(word, 0L) + count
          map1.update(word, newCount)
        }
      }
    }

    override def value: mutable.Map[String, Long] = {
      _count
    }
  }
}
